from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# 示例文档嵌入向量（通常通过预训练模型如BERT获取）
documents_embeddings = np.array([
    [0.1, 0.2, 0.3],
    [0.4, 0.5, 0.6],
    [0.7, 0.8, 0.9]
])

# 用户查询的嵌入向量
query_embedding = np.array([[0.2, 0.3, 0.4]])

# 计算查询嵌入与所有文档嵌入之间的余弦相似度
similarities = cosine_similarity(query_embedding, documents_embeddings)

# 找到最相似的文档索引
most_similar_doc_index = np.argmax(similarities)

print(f"最相似的文档索引: {most_similar_doc_index}")
print(f"余弦相似度: {similarities[0][most_similar_doc_index]}")